from pydantic.dataclasses import dataclass
from commandsniffer.commands.base import BaseCommand
from commandsniffer.graph import CommandGraph
from typing import Dict, List, Optional


@dataclass
class CommandModifyOperation:
    curr_command: BaseCommand

    def take_action(self, graph: CommandGraph) -> None:
        raise NotImplementedError


@dataclass
class ReplaceCommandOperation(CommandModifyOperation):
    new_commands: List[BaseCommand]

    def take_action(self, graph: CommandGraph) -> None:
        replace_commands: Dict[int, BaseCommand] = {}
        for cmd in self.new_commands:
            cmd.id = graph.new_command_id()
            replace_commands[cmd.id] = cmd
            graph.dependencies[cmd] = graph.dependencies.get(self.curr_command, set())
        graph.commands.update(replace_commands)
        graph.commands.pop(self.curr_command.id, self.curr_command)
        graph.dependencies.pop(self.curr_command, set())
        for dep in graph.dependencies.values():
            if self.curr_command in dep:
                dep.discard(self.curr_command)
                dep.update(replace_commands.values())


@dataclass
class AppendCommandBeforeOperation(CommandModifyOperation):
    added_commands: List[BaseCommand]

    def take_action(self, graph: CommandGraph) -> None:
        added_commands: Dict[int, BaseCommand] = {}
        curr_dep = graph.dependencies.get(self.curr_command, set())
        new_dep = set()
        for cmd in self.added_commands:
            cmd.id = graph.new_command_id()
            added_commands[cmd.id] = cmd
            graph.dependencies[cmd] = curr_dep
            new_dep.add(cmd)
        graph.commands.update(added_commands)
        graph.dependencies[self.curr_command] = new_dep


@dataclass
class AppendCommandAfterOperation(CommandModifyOperation):
    added_commands: List[BaseCommand]

    def take_action(self, graph: CommandGraph) -> None:
        added_commands = {}
        added_dep = set([self.curr_command])
        for cmd in self.added_commands:
            if cmd.id in graph.commands:
                cmd.id = graph.new_command_id()
            added_commands[cmd.id] = cmd
        graph.commands.update(added_commands)
        for dep in graph.dependencies.values():
            if self.curr_command in dep:
                dep.discard(self.curr_command)
                dep.update(added_commands.values())
        for cmd in added_commands.values():
            graph.dependencies[cmd] = added_dep


@dataclass
class RemoveCommandOperation(CommandModifyOperation):

    def take_action(self, graph: CommandGraph) -> None:
        graph.commands.pop(self.curr_command.id)
        graph.dependencies.pop(self.curr_command, set())
        for deps in graph.dependencies.values():
            deps.discard(self.curr_command)


@dataclass
class CommandPass:
    def run(self, command: BaseCommand) -> Optional[CommandModifyOperation]:
        raise NotImplementedError
